// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/operator_trace_writer.h>

#include <melon/hash/checksum.h>
#include <melon/io/cursor.h>
#include <utility>

#include <pollux/common/file/file.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/exec/operator.h>
#include <pollux/exec/trace.h>
#include <pollux/exec/trace_util.h>

namespace kumo::pollux::exec::trace {
    namespace {
        void recordOperatorSummary(Operator *op, melon::Dynamic &obj) {
            obj[OperatorTraceTraits::kOpTypeKey] = op->operatorType();
            const auto stats = op->stats(/*clear=*/false);
            if (op->operatorType() == "TableScan") {
                obj[OperatorTraceTraits::kNumSplitsKey] = stats.numSplits;
            }
            obj[OperatorTraceTraits::kPeakMemoryKey] =
                    stats.memoryStats.peakTotalMemoryReservation;
            obj[OperatorTraceTraits::kInputRowsKey] = stats.inputPositions;
            obj[OperatorTraceTraits::kInputBytesKey] = stats.inputBytes;
            obj[OperatorTraceTraits::kRawInputRowsKey] = stats.rawInputPositions;
            obj[OperatorTraceTraits::kRawInputBytesKey] = stats.rawInputBytes;
        }
    } // namespace

    OperatorTraceInputWriter::OperatorTraceInputWriter(
        Operator *traceOp,
        std::string traceDir,
        memory::MemoryPool *pool,
        UpdateAndCheckTraceLimitCB updateAndCheckTraceLimitCB)
        : traceOp_(traceOp),
          traceDir_(std::move(traceDir)),
          fs_(filesystems::getFileSystem(traceDir_, nullptr)),
          pool_(pool),
          serde_(getNamedVectorSerde(VectorSerde::Kind::kPresto)),
          updateAndCheckTraceLimitCB_(std::move(updateAndCheckTraceLimitCB)) {
        traceFile_ = fs_->openFileForWrite(getOpTraceInputFilePath(traceDir_));
        POLLUX_CHECK_NOT_NULL(traceFile_);
    }

    void OperatorTraceInputWriter::write(const RowVectorPtr &rows) {
        if (MELON_UNLIKELY(finished_)) {
            return;
        }

        if (batch_ == nullptr) {
            batch_ = std::make_unique<VectorStreamGroup>(pool_, serde_);
            batch_->createStreamTree(
                std::static_pointer_cast<const RowType>(rows->type()),
                1'000,
                &options_);
        }
        batch_->append(rows);

        // Serialize and write out each batch.
        IOBufOutputStream out(
            *pool_, nullptr, std::max<int64_t>(64 * 1024, batch_->size()));
        batch_->flush(&out);
        batch_->clear();
        auto iobuf = out.getIOBuf();
        updateAndCheckTraceLimitCB_(iobuf->computeChainDataLength());
        traceFile_->append(std::move(iobuf));
    }

    void OperatorTraceInputWriter::finish() {
        if (finished_) {
            return;
        }

        POLLUX_CHECK_NOT_NULL(
            traceFile_, "The query data writer has already been finished");
        traceFile_->close();
        traceFile_.reset();
        batch_.reset();
        writeSummary();
        finished_ = true;
    }

    void OperatorTraceInputWriter::writeSummary() const {
        const auto summaryFilePath = getOpTraceSummaryFilePath(traceDir_);
        const auto file = fs_->openFileForWrite(summaryFilePath);
        melon::Dynamic obj = melon::Dynamic::object;
        recordOperatorSummary(traceOp_, obj);
        file->append(melon::toJson(obj));
        file->close();
    }

    OperatorTraceSplitWriter::OperatorTraceSplitWriter(
        Operator *traceOp,
        std::string traceDir)
        : traceOp_(traceOp),
          traceDir_(std::move(traceDir)),
          fs_(filesystems::getFileSystem(traceDir_, nullptr)),
          splitFile_(fs_->openFileForWrite(getOpTraceSplitFilePath(traceDir_))) {
        POLLUX_CHECK_NOT_NULL(splitFile_);
    }

    void OperatorTraceSplitWriter::write(const exec::Split &split) const {
        // TODO: Supports group later once we have driver id mapping in trace node.
        POLLUX_CHECK(!split.hasGroup(), "Do not support grouped execution");
        POLLUX_CHECK(split.hasConnectorSplit());
        const auto splitObj = split.connectorSplit->serialize();
        const auto splitJson = melon::toJson(splitObj);
        auto ioBuf = serialize(splitJson);
        splitFile_->append(std::move(ioBuf));
    }

    void OperatorTraceSplitWriter::finish() {
        if (finished_) {
            return;
        }

        POLLUX_CHECK_NOT_NULL(
            splitFile_, "The query data writer has already been finished");
        splitFile_->close();
        writeSummary();
        finished_ = true;
    }

    // static
    std::unique_ptr<melon::IOBuf> OperatorTraceSplitWriter::serialize(
        const std::string &split) {
        const uint32_t length = split.length();
        const uint32_t crc32 = melon::crc32(
            reinterpret_cast<const uint8_t *>(split.data()), split.size());
        auto ioBuf =
                melon::IOBuf::create(sizeof(length) + split.size() + sizeof(crc32));
        melon::io::Appender appender(ioBuf.get(), 0);
        appender.writeLE(length);
        appender.push(reinterpret_cast<const uint8_t *>(split.data()), length);
        appender.writeLE(crc32);
        return ioBuf;
    }

    void OperatorTraceSplitWriter::writeSummary() const {
        const auto summaryFilePath = getOpTraceSummaryFilePath(traceDir_);
        const auto file = fs_->openFileForWrite(summaryFilePath);
        melon::Dynamic obj = melon::Dynamic::object;
        recordOperatorSummary(traceOp_, obj);
        file->append(melon::toJson(obj));
        file->close();
    }
} // namespace kumo::pollux::exec::trace
